import pyodbc
import pandas as pd
import os
from dotenv import load_dotenv
from logger import logging

load_dotenv()

class SQLServer:
    def __init__(self, database=None):
        """
        初始化数据库连接
        从环境变量读取连接信息,支持自定义数据库名
        """
        self.conn_str = (
            "DRIVER={ODBC Driver 17 for SQL Server};"
            f"SERVER={os.getenv('MSSQL_SERVER', '10.10.10.245')};"
            f"DATABASE={database or os.getenv('MSSQL_DATABASE', 'dz_his')};"
            f"UID={os.getenv('MSSQL_UID', 'sa')};"
            f"PWD={os.getenv('MSSQL_PWD', 'Aa123456')}"
        )
        self._conn = None

    @contextmanager
    def connection(self):
        """获取数据库连接的上下文管理器"""
        try:
            conn = pyodbc.connect(self.conn_str, timeout=5)
            yield conn
        except Exception as e:
            logging.error(f"数据库连接失败: {e}")
            raise
        finally:
            if conn:
                conn.close()

    def query_df(self, sql: str,index_col:Union[str,List[str]]=None, params: Union[tuple, dict] = None) -> Optional[pd.DataFrame]:
        """
        执行参数化查询并返回DataFrame
        :param sql: SQL查询语句(可包含?或:param占位符)
        :index_col: 指定返回DataFrame的索引列
        :param params: 参数元组或字典
        :return: DataFrame或None（查询失败或无数据时）
        """
        try:
            with self.connection() as conn:
                if params:
                    df = pd.read_sql(sql, conn, params=params)
                else:
                    df = pd.read_sql(sql, conn)
                if df.empty:
                    logging.warning(f"查询无数据: {sql}, params: {params}")
                    return None
                return df
        except Exception as e:
            logging.error(f"查询失败: {sql}\nparams: {params}\n错误: {e}")
            return None

    def query(self, sql: str, params: Union[tuple, dict] = None) -> List[tuple]:
        """
        执行参数化查询并返回结果列表
        :param sql: SQL查询语句(可包含?或:param占位符)
        :param params: 参数元组或字典
        :return: 结果列表，每行为一个元组
        """
        try:
            with self.connection() as conn:
                cursor = conn.cursor()
                if params:
                    cursor.execute(sql, params)
                else:
                    cursor.execute(sql)
                return cursor.fetchall()
        except Exception as e:
            logging.error(f"查询失败: {sql}\nparams: {params}\n错误: {e}")
            return []

    def execute(self, sql: str, params: Union[tuple, dict] = None) -> bool:
        """
        执行参数化SQL语句（INSERT/UPDATE/DELETE等）
        :param sql: SQL语句(可包含?或:param占位符)
        :param params: 参数元组或字典
        :return: 是否执行成功
        """
        try:
            with self.connection() as conn:
                cursor = conn.cursor()
                if params:
                    cursor.execute(sql, params)
                else:
                    cursor.execute(sql)
                conn.commit()
                return True
        except Exception as e:
            logging.error(f"执行失败: {sql}\nparams: {params}\n错误: {e}")
            return False

    def executemany(self, sql: str, params_list: List[Union[tuple, dict]]) -> bool:
        """
        批量执行参数化SQL语句
        :param sql: SQL语句(可包含?或:param占位符)
        :param params_list: 参数列表，每个元素为一组参数(元组或字典)
        :return: 是否执行成功
        """
        try:
            with self.connection() as conn:
                cursor = conn.cursor()
                cursor.executemany(sql, params_list)
                conn.commit()
                return True
        except Exception as e:
            logging.error(f"批量执行失败: {sql}\nparams数量: {len(params_list)}\n错误: {e}")
            return False

# 为了兼容性保留的函数
def query_sqlserver_df(sql: str, database: str = "dz_his") -> Optional[pd.DataFrame]:
    """
    兼容性函数：查询SQL Server并返回DataFrame
    建议直接使用 SQLServer 类
    """
    db = SQLServer(database=database)
    return db.query_df(sql)